import socket
import threading
import mxnet as mx
import logging
import numpy as np
from skimage import io, transform

logger = logging.getLogger()
logger.setLevel(logging.DEBUG)

prefix = "lenet"
num_round = 20
model = mx.model.FeedForward.load(prefix, num_round, ctx=mx.cpu(), numpy_batch_size=1)
mean_img = mx.nd.load("/home/wsy/mxnetMount/example/HCL2000/32data/mean.bin")["mean_img"]

synset = [l.strip() for l in open('synset2.txt').readlines()]
batch_size=1
data_shape = (1, 32, 32)


def PreprocessImage(path, show_img=False):
    # load image
    img = io.imread(path)
    # we crop image from center
    short_egde = min(img.shape[:2])
    yy = int((img.shape[0] - short_egde) / 2)
    xx = int((img.shape[1] - short_egde) / 2)
    crop_img = img[yy : yy + short_egde, xx : xx + short_egde]
    # resize
    #resized_img = transform.resize(crop_img, (32, 32))
    resized_img = crop_img[:,:,0]
    resized_img = transform.resize(resized_img, (32, 32))
    if show_img:
        io.imshow(resized_img)
    # convert to numpy.ndarray
    sample = np.asarray(resized_img) * 256
    # swap axes to make image from (224, 224, 4) to (3, 224, 224)
    #sample = np.swapaxes(sample, 0, 2)
    #sample = np.swapaxes(sample, 1, 2)
    # sub mean
    normed_img = sample - mean_img.asnumpy()
    normed_img.resize(1, 32, 32)
    normed_img=normed_img.reshape(1,1, 32, 32)
    return normed_img

class myThread (threading.Thread):
    def __init__(self, clientsocket):
        threading.Thread.__init__(self)
        self.clientsocket = clientsocket
    def run(self):
        try:
            self.clientsocket.settimeout(5)

            batch = PreprocessImage('1.jpg', True)
            prediction = model.predict(batch)[0]
            # print prediction.size #3755
            #print prediction.argmax()
            label = synset[prediction.argmax()]
            res = 'the result is ' + str(label) + ' and the probability is %f' % prediction.max()
            #print res

            self.clientsocket.send(res)

        except socket.timeout:
            print 'time out'
        self.clientsocket.close()

if __name__ == '__main__':
    host = "10.108.166.64"
    port = 8080
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.bind((host, port))
    sock.listen(500)
    print "Server is running on port %d" % port
    while True:
        clientsocket, address = sock.accept()
        #print "accept client: " + str(address)
        t = myThread(clientsocket)
        t.start()

